#!/usr/bin/env python

import struct
import os

from datetime import datetime, timedelta

import FATLFN

class FATDE:
   #             LE, 11 ASCII, Byte * 3, Int, Half * 2, Int, Half, Int
   FATDEFormat = "<11sBBBIHHIHI"
   def size(self):
      return struct.calcsize(self.FATDEFormat)

   def getDate(self, bits):
      # extract day, month year from bits
      day, month, year = bits & 0x1f, (bits >> 5) & 0xf, (bits >> 9) + 1980

      try:
         dt = datetime(year, month, day)
      except ValueError:
         # We will know it's an error cos it's not a datetime.
         dt = [year, month, day]

      return dt

   def getDateTime(self, bits):
      hi = bits >> 16
      lo = bits & 0xffff

      day, month, year = hi & 0x1f, (hi >> 5) & 0xf, (hi >> 9) + 1980

      hours, minutes, seconds = lo >> 11, (lo >> 5) & 0x3f, (lo & 0x1f) << 1

      try:
         dt = datetime(year, month, day, hours, minutes, seconds)
      except ValueError:
         dt = [year, month, day, hours, minutes, seconds]

      return dt
      
   def stringDate(self, d, micros=False):
      if isinstance(d, datetime):
         format = "%Y:%m:%d"
         if not(d.hour == d.minute == d.second == 0):
            format += "-%H:%M:%S"
         if micros:
            format += ".%f"

         string = d.strftime(format)
         if micros:
            string = string[0:-4]
         return string
      else:
         # For debugging ... just want to see what we tried to set the date to.
         return str(d)

   def longFileName(self):
      #longFileName = (self.attr == 0x0f)
      longFileName = (self.attr == 0x0f) and (self.reserved == 0) and (self.loCluster == 0)

      return longFileName

   def isLongFileName(self):
      return self.longFileName()

   def isArchive(self):
      return self.attr & 0x20 == 0x20

   def isDirectory(self):
      return self.attr & 0x10 == 0x10

   def isVolume(self):
      return self.attr & 0x08 == 0x08

   def isSystem(self):
      return self.attr & 0x04 == 0x04

   def isHidden(self):
      return self.attr & 0x02 == 0x02

   def isReadonly(self):
      return self.attr & 0x01 == 0x01

   def attrStr(self):
      string = 'a' if self.isArchive() else '-'
      string += 'd' if self.isDirectory() else '-'
      string += 'v' if self.isVolume() else '-'
      string += 's' if self.isSystem() else '-'
      string += 'h' if self.isHidden() else '-'
      string += 'r' if self.isReadonly() else '-'
      return string
      
   def getSize(self):
      return self.size
      
   def getName(self):
      return self.name
      
   def getSizeInClusters(self, clusterSize):
      if self.isDirectory():
         return 1
      elif self.getSize() == 0:
         return 0
      else:
         return (self.size + clusterSize - 1) / clusterSize

   def isDot(self):
      return self.name == '.' + 10 * ' '
      
   def isDotDot(self):
      return self.name == '..' + 9 * ' '
      
   def isDotted(self):
      return self.isDot() or self.isDotDot()


   def chksum(self):
      # Long file names link to a checksum of the short file name
      # Here we calculate the short filename checksum (from Brian Carrier)
      chksum = 0
      for asc in self.name:
         # Rotate chksum
         chksum = (0x80 if chksum & 0x01 else 0) + (chksum >> 1)
         chksum = (chksum + ord(asc)) & 0xff

      return chksum
   
   #
   #  Given two directory clusters, say whether they could be a match
   #  Match:
   #     1. LFN from 1 matches LFN from other (chksum and sequence)  2 (can be used to hunt for match)
   #     2. LFN from 1 matches SFN from other (chksum and sequence)  2
   #     3. SFN from 1 matches SFN from 2 or LFN start of sequence (or deleted!)   0
   #     4. Deleted LFN followed by deleted entry (LFN or SFN)       1
   #           but deleted LFN followed by deleted LFN can have chksum match
   #
   #     Otherwise no match -1
   #
   def matchFollowing(self, other):
      if self.isLongFileName():
         if self.deleted() and other.deleted(): # case 4
            return 1
            
         lfn = FATLFN.FATLFN(self.buffer)
         if other.isLongFileName():
            olfn = FATLFN.FATLFN(other.buffer)
            # Case 1, strong check, chksum and sequence match
            if lfn.chksum == olfn.chksum and (lfn.sequence & ~ 0x40) - 1 == olfn.sequence: # Case 1
               return 2 # strong match
            else:
               return -1 # no match
         else:
            # This time we need to calculate the chksum
            if lfn.chksum == other.chksum() and lfn.sequence & ~ 0x40 == 1: # Case 2
               return 2
            else:
               return -1
      else:
         # SFN here, other should be neutral, i.e. LFN start or another SFN
            
         if other.isLongFileName():
            olfn = FATLFN.FATLFN(other.buffer)
            if olfn.sequence & 0x40 == 0x40: # Start of LFN sequence
               return 0
            else:
               return -1
         else:
            return 0

   def getDeletedFirstChar(self, longName):
      if self.name[0] == 0xe5:
         self.name[0] = longName[0].uppercase()

      return self.chksum()
      
   def deleted(self):
      return ord(self.name[0]) == 0xe5

   def parse(self, buffer):
      item = struct.unpack(self.FATDEFormat, buffer)

      index = 0;
      self.name = item[index]    # 11s
      index += 1
      self.attr = item[index]    # B
      index += 1
      self.reserved = item[index]   # B
      self.loCluster = item[8]   # offset of cluster
      if not self.longFileName():
         index += 1
         crtime = item[index]
         index += 1
         self.createtime = self.getDateTime(item[index])
         # Add in the create time (hundreths of a sec)
         if isinstance(self.createtime, datetime):
            self.createtime += timedelta(0, 0, 0, crtime * 10)
         else:
            # tuple, add the create time
            self.createtime.append(crtime)

         index += 1
         self.access = self.getDate(item[index])
         index += 1
         self.hiCluster = item[index]
         index += 1
         self.modtime = self.getDateTime(item[index])
         index += 1
         index += 1
         self.size = item[index]

   def __init__(self, buffer):
      # keep the buffer so that we can reparse as a long file name if necessary
      self.offset = -1
      self.buffer = buffer
      self.parse(buffer)
      self.offset = -1,-1
      
   def endOfDirList(self):
      return ord(self.name[0]) == 0
      
   def validDOSchar(self, asc):
      # From wiki, invalid chars are ...
      c = ord(asc)
      if c >= 0 and c <= 31:
         return False
      if ord('a') <= c and c <= ord('z'):
         return False
      if asc in '"*/:<>?\\|':
         return False
      if asc in '+,.;=[]':
         return False
      if c == 127:
         return False
         
      return True # anything else is allowed
      
   def validName(self, name):
      if name == "." + 10 * " " or name == ".." + 9 * " ":
         return True

      # check all but the first
      for c in name[1:]:
         if not self.validDOSchar(c):
            return False # Any non valid chars => not a valid name
            
      # Now, check the first
      if self.validDOSchar(name[0]):
         return True # First is OK => we're OK

      # finally, if the file was deleted, the first char can be either ...
      first = ord(name[0])
      return first == 5 or first == 0xe5

   def isValid(self):
      # reserved should be zero, but I have seen 0x10 and 0x18 (rarely). Even 0x08.
      #valid = (self.reserved == 0) or (self.reserved == 0x10) or (self.reserved == 0x18) or (self.reserved == 0x08)
      valid = (self.reserved | 0x18) == 0x18
      # rules are different for long file name
      self.inValidMessage = "" if valid else "Reserved Byte not zero; " + hex(self.reserved)

      if not self.longFileName():
         if not self.validName(self.name):
            valid = False
            self.inValidMessage += "Bad name; " + self.name
            
         if not isinstance(self.createtime, datetime) or not isinstance(self.access, datetime) or not isinstance(self.modtime, datetime):
            valid = False
            self.inValidMessage += "Bad Dates; "
         else:
            # All the times are valid, check they're not in the future.
            future = datetime.now() + timedelta(2) # two years in future
            if self.createtime > future or self.access > future or self.modtime > future:
               valid = False
               self.inValidMessage += "Dates in the future; "

         # Weirdness: Have encountered a non zero directory size when the file was deleted.
         if not self.deleted() and self.isDirectory() and self.size != 0:
            valid = False
            self.inValidMessage += "Non zero directory size; "
         # Top two bits of attribute should be zero.
         if self.attr & 0xc0 != 0:
            valid = False
            self.inValidMessage += "attribute hi bytes not zero; "

      return valid

   def cluster(self):
      return (self.hiCluster << 16) + self.loCluster
      
   def getClusterNum(self):
      return self.cluster()
      
   def getNextFreeCluster(self, clusterSize):
      return self.cluster() + self.getSizeInClusters(clusterSize)
      
   # The cluster which contains this entry
   def getHoldingCluster(self):
      return self.offset[0]

   def getCreateDate(self):
      return self.createtime

   def __str__(self) :
      string = ""
      if self.longFileName():
         # Messy!!!
         fatlfn= FATLFN.FATLFN(self.buffer)
         string += str(fatlfn)
         #string += "LFN: " + hex(ord(self.name[0]))
      else:
         string += self.attrStr()
         string += ", " + self.stringDate(self.createtime, True)
         string += ", " + self.stringDate(self.access)
         #string += ", " + hex(self.hiclust)
         string += ", " + self.stringDate(self.modtime)
         string += ", " + hex(self.cluster())
         string += ", " + hex(self.getNextFreeCluster(8 * 512))
         string += ", " + str(self.size)

         string += ", " + hex(self.offset[0]) + ":"  + hex(self.offset[1])

         string += ", " + self.name + " (" + hex(self.chksum())[2:] + ")"
         
         string += ", " + ("[D]" if self.deleted() else "[K]")

      if not self.isValid():
         string += ", [" + self.inValidMessage + "]"

      return string

def isDirectory(sector):
   # First entry must be valid
   full = True # Assume dir is full
   first = sector[0:0x20]
   fatde = FATDE(first)
   if not fatde.isValid():
      return False
      
   # Search for first not valid entry
   for offset in range(0x20, len(sector), 0x20):
      de = sector[offset:offset+0x20]
      fatde = FATDE(de)
      if not fatde.isValid():
         full = False
         break
         #if not reduce(lambda x,y: x|y, [ord(z) for z in de]) == 0:
            # find out why invalid (if not all zeros)
            #  print fatde

   # Found a non valid directory entry...
   # ... this cluster will only be valid if the remainder of dir is zeros. (FIND REF)
   if full:
      return True
   else:
      # Sector isn't full, should be terminated with zeros
      valid = reduce(lambda x,y: x|y, [ord(z) for z in sector[offset:-1]]) == 0
   
      return valid

#
#  Find all sectors that contain valid FAT directories
#
def findDirectorySectors(imageFile, debug = False):
   sectorSize = 512
   sectorNum = 0
   

   image = open(imageFile)
   image.seek(0)
   
   sectors = []

   totalsectors = os.path.getsize(imageFile) / 512
   if debug: print "Searching for directory sectors (" + str(totalsectors) + " to search)."
   if totalsectors > 0:
      while True:
         sector = image.read(sectorSize)
         if len(sector) <> sectorSize:
            # terminated
            break
            
         if isDirectory(sector):
            sectors.append(sectorNum) 
            
         sectorNum += 1
         if debug and sectorNum % 256 == 0:
            print "\t" + str(sectorNum * 100 / totalsectors) + "% complete, (" + str(sectorNum) + " sectors processed).\r",
                  
   if debug:
      print # Otherwise next item to be printed will overwrite % message)
   
   return sectors


def readSector(imageFile, sectorNum):
   sectorSize = 512

   image = open(imageFile)
   
   image.seek(sectorNum*512)

   sector = image.read(sectorSize)
   if len(sector) <> sectorSize:
      sector = None
   return sector

if __name__ == '__main__':
   import sys
   from optparse import OptionParser

   parser = OptionParser(usage = "usage: %prog [options] imageFile")
   parser.add_option("-s", "--skip", type="int", dest="skip", default=0, help="The offset where valid clusters start (in bytes)")
   parser.add_option("-c", "--cluster", type="int", dest="clusterSize", default = 2 ** 10, help="Cluster size (in bytes)")

   (options, args) = parser.parse_args()

   if len(args) == 0:
      parser.print_help()
      sys.exit()

   directorySectors = findDirectorySectors(args[0])
   
   print directorySectors
